"""
摘要生成模块
提供文本摘要和关键词提取功能
"""

import re
import math
import heapq
import jieba
import jieba.analyse
import nltk
from typing import List, Dict, Tuple, Optional, Union
from collections import Counter
from sklearn.feature_extraction.text import TfidfVectorizer
from nltk.tokenize import sent_tokenize, word_tokenize
from nltk.corpus import stopwords
from .config import get_config

# 确保NLTK数据已下载
try:
    nltk.data.find('tokenizers/punkt')
    nltk.data.find('corpora/stopwords')
except LookupError:
    nltk.download('punkt')
    nltk.download('stopwords')

class Summarizer:
    """摘要生成类"""
    
    def __init__(self):
        """初始化摘要生成类"""
        # 获取配置
        self.algorithm = get_config("摘要.算法", "extractive")
        self.default_ratio = get_config("摘要.默认比例", 0.3)
        self.keywords_count = get_config("摘要.关键词数量", 10)
        
        # 加载停用词
        self.stopwords_zh = self._load_stopwords_zh()
        self.stopwords_en = set(stopwords.words('english'))
    
    def _load_stopwords_zh(self) -> set:
        """
        加载中文停用词
        
        Returns:
            set: 中文停用词集合
        """
        # 常用中文停用词
        common_stopwords = {
            '的', '了', '和', '是', '就', '都', '而', '及', '与', '着',
            '或', '一个', '没有', '我们', '你们', '他们', '她们', '它们',
            '这个', '那个', '这些', '那些', '这样', '那样', '之', '的话',
            '一下', '一种', '一样', '一起', '不过', '不要', '这', '那',
            '才', '也', '么', '却', '还', '只', '但', '但是', '何', '啊'
        }
        return common_stopwords
    
    def _is_chinese(self, text: str) -> bool:
        """
        判断文本是否主要为中文
        
        Args:
            text: 输入文本
            
        Returns:
            bool: 是否主要为中文
        """
        # 统计中文字符数量
        chinese_chars = sum(1 for char in text if '\u4e00' <= char <= '\u9fff')
        # 如果中文字符占比超过20%，则认为是中文文本
        return chinese_chars / max(len(text), 1) > 0.2
    
    def _split_sentences(self, text: str) -> List[str]:
        """
        分割文本为句子
        
        Args:
            text: 输入文本
            
        Returns:
            List[str]: 句子列表
        """
        # 判断文本语言
        is_chinese = self._is_chinese(text)
        
        if is_chinese:
            # 中文分句
            text = re.sub(r'([。！？\?!])+', r'\1\n', text)
            return [s.strip() for s in text.split('\n') if s.strip()]
        else:
            # 英文分句
            return sent_tokenize(text)
    
    def _tokenize(self, text: str, is_chinese: bool) -> List[str]:
        """
        分词
        
        Args:
            text: 输入文本
            is_chinese: 是否为中文
            
        Returns:
            List[str]: 词语列表
        """
        if is_chinese:
            # 中文分词
            return [word for word in jieba.cut(text) if word.strip()]
        else:
            # 英文分词
            return word_tokenize(text.lower())
    
    def _remove_stopwords(self, words: List[str], is_chinese: bool) -> List[str]:
        """
        移除停用词
        
        Args:
            words: 词语列表
            is_chinese: 是否为中文
            
        Returns:
            List[str]: 移除停用词后的词语列表
        """
        if is_chinese:
            return [word for word in words if word not in self.stopwords_zh]
        else:
            return [word for word in words if word not in self.stopwords_en]
    
    def _calculate_word_frequencies(self, words: List[str]) -> Dict[str, int]:
        """
        计算词频
        
        Args:
            words: 词语列表
            
        Returns:
            Dict[str, int]: 词频字典
        """
        return Counter(words)
    
    def _calculate_sentence_scores(self, sentences: List[str], word_freq: Dict[str, int], is_chinese: bool) -> Dict[int, float]:
        """
        计算句子得分
        
        Args:
            sentences: 句子列表
            word_freq: 词频字典
            is_chinese: 是否为中文
            
        Returns:
            Dict[int, float]: 句子索引到得分的映射
        """
        sentence_scores = {}
        
        for i, sentence in enumerate(sentences):
            # 分词
            words = self._tokenize(sentence, is_chinese)
            # 移除停用词
            words = self._remove_stopwords(words, is_chinese)
            
            # 计算句子得分
            score = 0
            for word in words:
                if word in word_freq:
                    score += word_freq[word]
            
            # 归一化得分（按句子长度）
            if len(words) > 0:
                score = score / len(words)
            
            sentence_scores[i] = score
        
        return sentence_scores
    
    def _extractive_summarize(self, text: str, ratio: float = None) -> str:
        """
        抽取式摘要
        
        Args:
            text: 输入文本
            ratio: 摘要比例，默认使用配置值
            
        Returns:
            str: 摘要文本
        """
        if ratio is None:
            ratio = self.default_ratio
        
        # 分割句子
        sentences = self._split_sentences(text)
        
        # 如果句子数量太少，直接返回原文
        if len(sentences) <= 3:
            return text
        
        # 判断文本语言
        is_chinese = self._is_chinese(text)
        
        # 将所有句子合并为一个文本，用于分词
        all_text = ' '.join(sentences)
        
        # 分词
        words = self._tokenize(all_text, is_chinese)
        
        # 移除停用词
        words = self._remove_stopwords(words, is_chinese)
        
        # 计算词频
        word_freq = self._calculate_word_frequencies(words)
        
        # 计算句子得分
        sentence_scores = self._calculate_sentence_scores(sentences, word_freq, is_chinese)
        
        # 选择得分最高的句子
        summary_count = max(1, int(len(sentences) * ratio))
        top_indices = heapq.nlargest(summary_count, sentence_scores, key=sentence_scores.get)
        
        # 按原始顺序排列句子
        top_indices.sort()
        
        # 生成摘要
        summary = [sentences[i] for i in top_indices]
        
        return ' '.join(summary)
    
    def _abstractive_summarize(self, text: str, ratio: float = None) -> str:
        """
        生成式摘要（简化版，实际应用中可以使用更复杂的模型）
        
        Args:
            text: 输入文本
            ratio: 摘要比例，默认使用配置值
            
        Returns:
            str: 摘要文本
        """
        # 由于生成式摘要通常需要复杂的模型，这里简化为抽取式摘要的变体
        # 在实际应用中，可以集成预训练的语言模型如BART、T5等
        
        # 分割句子
        sentences = self._split_sentences(text)
        
        # 如果句子数量太少，直接返回原文
        if len(sentences) <= 3:
            return text
        
        # 判断文本语言
        is_chinese = self._is_chinese(text)
        
        # 使用TF-IDF向量化句子
        if is_chinese:
            # 中文需要先分词
            tokenized_sentences = [' '.join(jieba.cut(s)) for s in sentences]
            vectorizer = TfidfVectorizer(stop_words=list(self.stopwords_zh))
        else:
            tokenized_sentences = sentences
            vectorizer = TfidfVectorizer(stop_words='english')
        
        # 计算TF-IDF矩阵
        tfidf_matrix = vectorizer.fit_transform(tokenized_sentences)
        
        # 计算句子得分（使用TF-IDF值的和）
        sentence_scores = {}
        for i, sentence in enumerate(sentences):
            sentence_scores[i] = sum(tfidf_matrix[i].toarray()[0])
        
        # 选择得分最高的句子
        if ratio is None:
            ratio = self.default_ratio
        summary_count = max(1, int(len(sentences) * ratio))
        top_indices = heapq.nlargest(summary_count, sentence_scores, key=sentence_scores.get)
        
        # 按原始顺序排列句子
        top_indices.sort()
        
        # 生成摘要
        summary = [sentences[i] for i in top_indices]
        
        return ' '.join(summary)
    
    def summarize(self, text: str, ratio: float = None, algorithm: str = None) -> str:
        """
        生成文本摘要
        
        Args:
            text: 输入文本
            ratio: 摘要比例，默认使用配置值
            algorithm: 摘要算法，默认使用配置值
            
        Returns:
            str: 摘要文本
        """
        if not text:
            return ""
        
        if algorithm is None:
            algorithm = self.algorithm
        
        if ratio is None:
            ratio = self.default_ratio
        
        # 根据算法选择相应的摘要方法
        if algorithm == "extractive":
            return self._extractive_summarize(text, ratio)
        elif algorithm == "abstractive":
            return self._abstractive_summarize(text, ratio)
        else:
            # 默认使用抽取式摘要
            return self._extractive_summarize(text, ratio)
    
    def extract_keywords(self, text: str, count: int = None) -> List[str]:
        """
        提取关键词
        
        Args:
            text: 输入文本
            count: 关键词数量，默认使用配置值
            
        Returns:
            List[str]: 关键词列表
        """
        if not text:
            return []
        
        if count is None:
            count = self.keywords_count
        
        # 判断文本语言
        is_chinese = self._is_chinese(text)
        
        if is_chinese:
            # 使用jieba提取中文关键词
            keywords = jieba.analyse.extract_tags(text, topK=count)
        else:
            # 英文关键词提取
            # 分词
            words = word_tokenize(text.lower())
            # 移除停用词
            words = [word for word in words if word.isalnum() and word not in self.stopwords_en]
            # 计算词频
            word_freq = Counter(words)
            # 选择频率最高的词作为关键词
            keywords = [word for word, _ in word_freq.most_common(count)]
        
        return keywords
    
    def set_algorithm(self, algorithm: str) -> None:
        """
        设置摘要算法
        
        Args:
            algorithm: 摘要算法（"extractive"或"abstractive"）
        """
        if algorithm in ["extractive", "abstractive"]:
            self.algorithm = algorithm
    
    def set_default_ratio(self, ratio: float) -> None:
        """
        设置默认摘要比例
        
        Args:
            ratio: 摘要比例（0-1之间）
        """
        if 0 < ratio <= 1:
            self.default_ratio = ratio
    
    def set_keywords_count(self, count: int) -> None:
        """
        设置关键词数量
        
        Args:
            count: 关键词数量
        """
        if count > 0:
            self.keywords_count = count

# 创建全局实例
summarizer = Summarizer()

def init_summarizer() -> Summarizer:
    """
    初始化全局摘要生成实例
    
    Returns:
        Summarizer: 摘要生成实例
    """
    global summarizer
    summarizer = Summarizer()
    return summarizer

def summarize(text: str, ratio: float = None, algorithm: str = None) -> str:
    """
    生成文本摘要的快捷方法
    
    Args:
        text: 输入文本
        ratio: 摘要比例
        algorithm: 摘要算法
        
    Returns:
        str: 摘要文本
    """
    return summarizer.summarize(text, ratio, algorithm)

def extract_keywords(text: str, count: int = None) -> List[str]:
    """
    提取关键词的快捷方法
    
    Args:
        text: 输入文本
        count: 关键词数量
        
    Returns:
        List[str]: 关键词列表
    """
    return summarizer.extract_keywords(text, count)